from . import custommodule
from datetime import datetime,timedelta
import sqlite3


#全量更新
def full_update():
    """
    全量更新：
    1. 更新股票列表
    2. 更新所有股票的日线数据
    3. 重新计算所有技术指标
    """

    # 检查是否需要执行全量更新
    config = custommodule.ConfigManager() # 创建配置管理器
    start_date = datetime.strptime(config.get_value('full_down_time'), "%Y-%m-%d %H:%M:%S")
    if start_date.hour > 18:
        start_date = (start_date.date() + timedelta(days=1))
    start_date = start_date.strftime("%Y-%m-%d")
    end_date = datetime.now().date().strftime("%Y-%m-%d")
    trade_dates = custommodule.TradeDate.get_trade_dates(start_date, end_date)
    if trade_dates.empty:
        print("没有可更新的交易日数据！")
    else:
        if (datetime.now().replace(hour=18, minute=0, second=0, microsecond=0) - timedelta(days=1) < datetime.strptime(
                config.get_value('full_down_time'),
                "%Y-%m-%d %H:%M:%S") and datetime.now().hour < 18) or datetime.strptime(
                config.get_value('full_down_time'), "%Y-%m-%d %H:%M:%S") > datetime.now().replace(hour=18, minute=0, second=0, microsecond=0):
            choice = input(f"最近全量下载日期为{config.get_value('full_down_time')}，是否继续执行全量下载？（y/n）")
            if choice.lower() != 'y':
                print("已取消全量下载")
            else:
                print("🔄 开始执行全量下载...")
                custommodule.update_stock_list()  # 更新股票列表
                stock_codes = [code for code, _ in custommodule.get_all_stocks_info()]  # 获取所有股票代码
                custommodule.update_stock_daily(stock_codes)  # 下载所有股票的全日期数据
                config.set_value('full_down_time', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))  # 保存下载时间
                print('全量下载完成')
        else:
            print("🔄 开始执行全量下载...")
            custommodule.update_stock_list()  # 更新股票列表
            stock_codes = [code for code, _ in custommodule.get_all_stocks_info()]  # 获取所有股票代码
            custommodule.update_stock_daily(stock_codes)  # 下载所有股票的全日期数据
            config.set_value('full_down_time', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))  # 保存下载时间
            print('全量下载完成')

    # 检查是否需要执行全量技术指标计算
    if datetime.strptime(config.get_value('full_down_time'), "%Y-%m-%d %H:%M:%S").date() > datetime.strptime(config.get_value('full_calculate_time'), "%Y-%m-%d %H:%M:%S").date():
        print("🔄 开始执行全量技术指标计算...")
        stock_codes = [code for code, _ in custommodule.get_all_stocks_info()]  # 获取所有股票代码
        custommodule.calculate_technical_indicators(stock_codes)  # 计算所有股票的技术指标
        config.set_value('full_calculate_time', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))  # 保存计算时间
        print('全量计算完成')
    else:
        print('无需计算技术指标')

#增量更新
def incremental_update():
    """
    1. 下载更新股票列表
    2. 将所有股票分成两类：最近复权、最近非复权
    3. 对最近复权的股票进行全量更新
    4. 对最近非复权的股票进行增量更新
    """
    config = custommodule.ConfigManager()  # 创建配置管理器
    # 如果全量计算时间大于增量更新时间，则将增量更新时间设置为全量计算时间
    if config.get_value('incremental_update_time') < config.get_value('full_calculate_time'):
        config.set_value('incremental_update_time', config.get_value('full_calculate_time'))

    # 如果在增量更新时间到今天这个时间段内没有交易日，则无需更新
    start_date = datetime.strptime(config.get_value('incremental_update_time'), "%Y-%m-%d %H:%M:%S")
    if start_date.hour > 18:
        start_date = start_date.date() + timedelta(days=1)
    start_date = start_date.strftime("%Y-%m-%d")
    end_date = datetime.now().date().strftime("%Y-%m-%d")
    trade_dates = custommodule.TradeDate.get_trade_dates(start_date, end_date)
    if trade_dates.empty:
        print("没有可更新的交易日数据！")

    # 若满足以下任一条件则不更新：
        # 上次增量更新时间在昨天 18 点之后，且当前时间在 18 点之前。
        # 上次增量更新时间在今天 18 点之后。
    # 若不满足上述条件，则开始执行增量更新。
    if (datetime.now().replace(hour=18, minute=0, second=0, microsecond=0) - timedelta(days=1) < datetime.strptime( config.get_value('incremental_update_time'), "%Y-%m-%d %H:%M:%S") and datetime.now().hour < 18) or datetime.strptime( config.get_value('incremental_update_time'), "%Y-%m-%d %H:%M:%S") > datetime.now().replace(hour=18, minute=0, second=0, microsecond=0):
        choice = input(f"最近增量更新日期为{config.get_value('incremental_update_time')}，是否继续执行增量更新？（y/n）")
        if choice.lower() != 'y':
            print("已取消增量更新")
        else:
            print("🔄 开始执行增量更新...")
            try:
                custommodule.update_stock_list()  # 更新股票列表

                # 根据股票名称将股票分为两类：应复权、非复权两类，应复权全量更新、非复权增量更新
                with sqlite3.connect('db_file.db') as conn:
                    cursor = conn.cursor()

                    # 获取应复权股列表
                    cursor.execute("""
                                SELECT stock_code, stock_name
                                FROM stock_info
                                WHERE stock_name LIKE 'XD%'
                                  OR stock_name LIKE 'XR%'
                                  OR stock_name LIKE 'DR%'
                            """)
                    adjusted_stocks = [code for code, _ in cursor.fetchall()]
                    # 获取非复权股列表
                    no_adjusted_stocks = [code for code, _ in cursor.execute("""
                                SELECT stock_code, stock_name 
                                FROM stock_info 
                                WHERE stock_name NOT LIKE 'XD%'
                                  AND stock_name NOT LIKE 'XR%'
                                  AND stock_name NOT LIKE 'DR%'
                            """)]

                    # 下载应复权股数据
                    if adjusted_stocks:
                        print("🔄 对应复权股票的日线数据进行全量更新...")
                        custommodule.update_stock_daily(adjusted_stocks)
                        print("🔄 对应复权股票的技术指标进行全量计算...")
                        custommodule.calculate_technical_indicators(adjusted_stocks)
                    else:
                        print("没有应复权的股票")
                    # 下载非复权股数据
                    if no_adjusted_stocks:
                        start_date = config.get_value('incremental_update_time')[:10]  # 获取上次更新时间
                        print("🔄 对非复权的股票进行增量更新...")
                        custommodule.update_stock_daily(no_adjusted_stocks, start_date)
                        # 保存下载时间
                        config.set_value('incremental_update_time', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
                        print("🔄 对非复权股票的技术指标进行增量计算...")
                        custommodule.calculate_technical_indicators(no_adjusted_stocks, start_date)
                    else:
                        print("没有非复权的股票")

                    print("✅ 增量更新完成")
            except sqlite3.Error as e:
                print(f"数据库操作失败: {e}")
            except Exception as e:
                print(f"未知错误发生: {e}")
    else:
        print("🔄 开始执行增量更新...")
        try:
            custommodule.update_stock_list()  # 更新股票列表

            # 根据股票名称将股票分为两类：应复权、非复权两类，应复权全量更新、非复权增量更新
            with sqlite3.connect('db_file.db') as conn:
                cursor = conn.cursor()

                # 获取应复权股列表
                cursor.execute("""
                    SELECT stock_code, stock_name
                    FROM stock_info
                    WHERE stock_name LIKE 'XD%'
                      OR stock_name LIKE 'XR%'
                      OR stock_name LIKE 'DR%'
                """)
                adjusted_stocks = [code for code, _ in cursor.fetchall()]
                # 获取非复权股列表
                no_adjusted_stocks = [code for code, _ in cursor.execute("""
                    SELECT stock_code, stock_name 
                    FROM stock_info 
                    WHERE stock_name NOT LIKE 'XD%'
                      AND stock_name NOT LIKE 'XR%'
                      AND stock_name NOT LIKE 'DR%'
                """)]

                # 下载应复权股数据
                if adjusted_stocks:
                    print("🔄 对应复权股票的日线数据进行全量更新...")
                    custommodule.update_stock_daily(adjusted_stocks)
                    print("🔄 对应复权股票的技术指标进行全量计算...")
                    custommodule.calculate_technical_indicators(adjusted_stocks)
                else:
                    print("没有应复权的股票")
                # 下载非复权股数据
                if no_adjusted_stocks:
                    start_date = config.get_value('incremental_update_time')[:10] #获取上次更新时间
                    print("🔄 对非复权的股票进行增量更新...")
                    custommodule.update_stock_daily(no_adjusted_stocks, start_date)
                    # 保存下载时间
                    config.set_value('incremental_update_time', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
                    print("🔄 对非复权股票的技术指标进行增量计算...")
                    custommodule.calculate_technical_indicators(no_adjusted_stocks, start_date)
                else:
                    print("没有非复权的股票")

                print("✅ 增量更新完成")
        except sqlite3.Error as e:
            print(f"数据库操作失败: {e}")
        except Exception as e:
            print(f"未知错误发生: {e}")


#自选股全量更新
def watchlist_update_1():
    """
    1. 获取自选股列表
    2. 更新自选股的所有日线数据
    3. 重新计算自选股的所有技术指标
    """
    config = custommodule.ConfigManager()  # 创建配置管理器

    # 如果全量计算时间大于自选股更新时间，则将自选股更新时间设置为全量计算时间
    if config.get_value('watchlist_update_time1') < config.get_value('full_calculate_time'):
        config.set_value('watchlist_update_time1', config.get_value('full_calculate_time'))

    # 如果自选股全量更新时间在昨天 18 点之后，且当前时间在 18 点之前，或者自选股全量更新时间在今天 18 点之后，则无需更新
    if (datetime.now().replace(hour=18, minute=0, second=0, microsecond=0) - timedelta(days=1) < datetime.strptime( config.get_value('watchlist_update_time1'), "%Y-%m-%d %H:%M:%S") and datetime.now().hour < 18) or datetime.strptime( config.get_value('watchlist_update_time1'), "%Y-%m-%d %H:%M:%S") > datetime.now().replace(hour=18, minute=0, second=0, microsecond=0):
        print("当前无须更新")
    else:
        stock_codes = [code for code, _ in custommodule.get_watchlist()]  # 获取自选股代码列表
        if stock_codes:
            print("🔄 开始执行自选股更新...")
            custommodule.update_stock_daily(stock_codes)  # 全量更新自选股日线数据
            config.set_value('watchlist_update_time1', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            custommodule.calculate_technical_indicators(stock_codes)  # 全量计算自选股技术指标
            print("✅ 自选股更新完成")
        else:
            print("⚠️ 自选股列表为空")


#自选股增量更新
def watchlist_update_2():
    """
    1. 获取自选股列表
    2. 更新自选股的所有日线数据
    3. 重新计算自选股的所有技术指标
    """
    config = custommodule.ConfigManager()  # 创建配置管理器

    # 如果自选股增量更新时间小于全量更新时间、自选股全量更新时间、自选股增量更新时间，则将自选股增量更新时间设置为三者中的最大值
    if config.get_value('watchlist_update_time2') < max(config.get_value('watchlist_update_time1'), config.get_value('full_calculate_time'), config.get_value('incremental_update_time')):
        config.set_value('watchlist_update_time2', max(config.get_value('watchlist_update_time1'), config.get_value('full_calculate_time'), config.get_value('incremental_update_time')))

    # 如果自选股增量更新时间在昨天 18 点之后，且当前时间在 18 点之前，或者自选股增量更新时间在今天 18 点之后，则无需更新
    if (datetime.now().replace(hour=18, minute=0, second=0, microsecond=0) - timedelta(days=1) < datetime.strptime(config.get_value('watchlist_update_time2'), "%Y-%m-%d %H:%M:%S") and datetime.now().hour < 18) or datetime.strptime(config.get_value('watchlist_update_time2'), "%Y-%m-%d %H:%M:%S") > datetime.now().replace(hour=18, minute=0, second=0, microsecond=0):
        print("自选股当前无须增量更新")
    else:
        stock_date = config.get_value('watchlist_update_time2')[:10]  # 获取自选股上次更新时间
        stock_codes = [code for code, _ in custommodule.get_watchlist()]  # 获取自选股代码列表
        if stock_codes:
            print("🔄 开始执行自选股增量更新...")
            custommodule.update_stock_daily(stock_codes, stock_date)  # 增量更新自选股日线数据
            config.set_value('watchlist_update_time2', datetime.now().strftime("%Y-%m-%d %H:%M:%S"))
            custommodule.calculate_technical_indicators(stock_codes, stock_date)  # 增量计算自选股技术指标
            print("✅ 自选股增量更新完成")
        else:
            print("⚠️ 自选股列表为空")


# 手动更新下载失败的股票
def error_update():
    config = custommodule.ConfigManager()
    stock_codes = config.get_value('error_codes', [])
    config.set_value('error_codes', "")
    if stock_codes:
        print("🔄 开始执行手动更新下载失败的股票...")
        stock_codes = stock_codes.split(",")
        custommodule.update_stock_daily(stock_codes)
        custommodule.calculate_technical_indicators(stock_codes)
        print("✅ 失败股更新完成")
    else:
        print("没有下载失败的股票")

# 清理数据库中冗余数据
def db_clear():
    sql = f"""
    -- 第一步：禁用外键约束
    PRAGMA foreign_keys = OFF;
    -- 第二步：先删除子表(stock_technical)中不匹配的数据
    DELETE FROM stock_technical
    WHERE stock_code NOT IN (
    SELECT SUBSTR(stock_code, 1, LENGTH(stock_code) - 3) 
    FROM stock_info
    );
    -- 第三步：再删除父表(stock_daily)中不匹配的数据
    DELETE FROM stock_daily
    WHERE stock_code NOT IN (
    SELECT SUBSTR(stock_code, 1, LENGTH(stock_code) - 3) 
    FROM stock_info
    );
    -- 第四步：重新启用外键约束
    PRAGMA foreign_keys = ON;
    """
    try:
        print("开始清理数据库冗余数据...")
        with sqlite3.connect('db_file.db') as conn:
            cursor = conn.cursor()
            cursor.executescript(sql)
            conn.commit()
        print("数据库清理完成！")
    except sqlite3.Error as e:
        print(f"数据库清理失败: {e}")


# 更新模块的菜单
def menu():
    """股票数据更新主菜单"""
    while True:
        print("\n" + "=" * 20)
        print(" 股票数据更新菜单 ")
        print("=" * 20)
        print("1. 全量更新")
        print("2. 增量更新")
        print("3. 自选股全量更新")
        print("4. 自选股增量更新")
        print("5. 失败股手动更新")
        print("6. 数据库清理")
        print("0. 返回主菜单")
        print("-" * 20)
        choice = input("请输入选项（0-4）：").strip()

        # 处理用户选择
        match choice:
            case '1':
                full_update()
            case '2':
                incremental_update()
            case '3':
                watchlist_update_1()
            case '4':
                watchlist_update_2()
            case '5':
                error_update()
            case '6':
                db_clear()
            case '0':
                break
            case _:
                print("⚠️ 输入错误，请输入 0-5 之间的数字")
